# cython: profile=True
import cython
from numpy cimport ndarray
from v3dsolver.fluid_elements.fluid_elems cimport FluidElem,Saffman,PseudoParticle
from v3dsolver.vector cimport Vec
from mpi4py cimport MPI

cdef class Cell:
    cdef public int level
    cdef public Vec center
    cdef public double length
    cdef public double radius
    cdef public int npmax
    cdef public double rad
    cdef public int maxlevel
    cdef public list parentelems
    cdef public list elems
    cdef public list cause_elems
    cdef public list effect_elems
    cdef public int npseudo
    cdef public list children
    cdef public bint childless
    cdef public int id

    @cython.locals(elems=list,elem=FluidElem)
    cdef list get_elems(self,int cause_or_effect=*)

    @cython.locals(elem_pos=Vec,self_center=Vec)
    cdef bint contains(self,FluidElem elem)

    @cython.locals(directions=list,corners=list,corner=ndarray)
    cdef list get_corners(self)

    @cython.locals(children=list,directions=list,direction=Vec,childcenter=Vec,child=Cell)
    cdef list get_children(self)

    cdef bint splittable(self)

    @cython.locals(vel=Vec,cause_elem=Saffman,v=Vec)
    cdef Vec direct_vel(self,FluidElem effect_elem,list cause_elems=*)

    @cython.locals(otherlen=double,otherrad=double,dist=Vec,distance=double)
    cdef bint is_neighbour(self,Vec cellcenter,int level)

    @cython.locals(dist=Vec,distance=double)
    cdef bint is_well_separated(self,FluidElem elem)

cdef class MakinoCell(Cell):
    cdef public MakinoCell parent
    cdef public list pseudo_p
    cdef public list P_l
    
    @cython.locals(children=list,directions=list,direction=Vec,childcenter=Vec,child=MakinoCell)
    cdef list get_children(self)

    @cython.locals(pseudo_p=list,newpseudo=PseudoParticle,pos=Vec,pseudo_pos_list=list)
    cdef list get_pseudo(self,list constituting_particles=*)

    @cython.locals(pseudo_elems=list,child=MakinoCell)
    cdef transfer_pseudo(self)

    @cython.locals(vel=Vec,child=MakinoCell,v=Vec)
    cdef Vec compute_vel(self,elem)

cdef class TreeManager:
    cdef public MPI.Comm comm
    cdef public int npmax
    cdef public list domain
    cdef public list elems
    cdef public int npseudo
    cdef public MakinoCell root
    cdef public list childless_cells
    cdef public double rad
    cdef public list cells
    cdef public int maxlevel
    cdef public list P_l
    
    @cython.locals(cell=MakinoCell,level=int)
    cdef generate_pseudo_particles(self)
    
    @cython.locals(pointer=int,level=int,index=int)
    cdef MakinoCell get_cell(self,int id)
